package com.zyk.grate_offer.class03;

import java.util.*;

/**
 * 给定三个参数：
 * 二叉树的头节点head，树上某个节点target，正数K
 * 从target开始，可以向上走或者向下走
 * 返回与target的距离是K的所有节点
 *
 * @author zhangsan
 * @date 2021/4/28 13:34
 */
public class Code08_DistanceKNodes {

    public static class Node {
        public int value;
        public Node left;
        public Node right;

        public Node(int v) {
            value = v;
        }
    }

    public static List<Node> distanceKNodes(Node root, Node target, int K) {
        List<Node> ans = new ArrayList<>();
        Map<Node, Node> parentMap = new HashMap<>();    // 可向上走，记录父节点好像上走
        parentMap.put(root, null);
        fillParentMap(root, parentMap);

        Queue<Node> q = new LinkedList<>();
        q.add(target);
        HashSet<Node> visited = new HashSet<>();
        visited.add(target);
        int curLevel = 0;
        Node cur, parent;
        while (!q.isEmpty()) {
            int size = q.size();
            while (size-- > 0) {
                cur = q.poll();
                if (cur.left != null && !visited.contains(cur.left)) {
                    q.add(cur.left);
                    visited.add(cur.left);
                }
                if (cur.right != null && !visited.contains(cur.right)) {
                    q.add(cur.right);
                    visited.add(cur.right);
                }
                if ((parent = parentMap.get(cur)) != null && !visited.contains(parent)) {
                    q.add(parent);
                    visited.add(parent);
                }
                if (curLevel == K) {
                    ans.add(cur);
                }
            }
            curLevel++;
            if (curLevel > K) {
                break;
            }
        }
        return ans;
    }

    private static void fillParentMap(Node root, Map<Node, Node> parentMap) {
        if (root == null) {
            return;
        }
        if (root.left != null) {
            parentMap.put(root.left, root);
            fillParentMap(root.left, parentMap);
        }
        if (root.right != null) {
            parentMap.put(root.right, root);
            fillParentMap(root.right, parentMap);
        }
    }

    public static void main(String[] args) {
        Node n0 = new Node(0);
        Node n1 = new Node(1);
        Node n2 = new Node(2);
        Node n3 = new Node(3);
        Node n4 = new Node(4);
        Node n5 = new Node(5);
        Node n6 = new Node(6);
        Node n7 = new Node(7);
        Node n8 = new Node(8);

        n3.left = n5;
        n3.right = n1;
        n5.left = n6;
        n5.right = n2;
        n1.left = n0;
        n1.right = n8;
        n2.left = n7;
        n2.right = n4;

        Node root = n3;
        Node target = n5;
        int K = 2;

        List<Node> ans = distanceKNodes(root, target, K);
        for (Node o1 : ans) {
            System.out.println(o1.value);
        }

    }


}
